import numpy.matlib as np
import infiniteSitesTools as ist

class SiteFrequencySpectrum:
	
	sfs = []
	#sfsNormed = []

	def __init__(self, polys=None):
		self.sfs = []
		self.sfsNormed = []
		if polys is not None:
			self.addSFSs(polys)

	def addSFS(self, inSFS):
		self.sfs.append(inSFS)
		
#		segSites = ist.segregatingSites(inSFS)
#		newSFS = inSFS / segSites
#		self.sfsNormed.append(newSFS)

	def addSFSs(self, polyTables):
		for p in polyTables:
			self.addSFS(self.getSiteFrequencySpectrum(p.table))

	def getExpectedSFS(self):
		data = np.matrix(self.sfs)
		numCols = data.shape[1]
		numRows = data.shape[0]

		sumCols = data.sum(axis=0)
		eSFS = sumCols / numRows

		return eSFS
	
#	def getExpectedSFSnormed(self):
#		data = np.matrix(self.sfsNormed)
#		numCols = data.shape[1]
#		numRows = data.shape[0]
#
#		sumCols = data.sum(axis=0)
#		eSFS = sumCols / numRows
#		
#		return eSFS
	
	def getTwoSidedSFS(self, p):
		data = np.matrix(self.sfsNormed)

		numRows = data.shape[0]
		lower = int(p * numRows)
		upper = numRows - lower
		
		sortedData = np.msort(data)
		
		sData = []
		sData.append(np.ravel(sortedData[lower]))
		sData.append(np.ravel(sortedData[upper]))
		
		rData = np.matrix(sData)
		return rData

	def getSiteFrequencySpectrum(self, polyTable):
		P = np.matrix(polyTable)
		n = P.shape[0]
		S = P.shape[1]
	
		sums = np.ravel(P.sum(axis=0))
		SFS = np.ravel(np.zeros(n, int))
	
		for cItr in range(S):
			k = sums[cItr]
			SFS[k] = SFS[k] + 1
		
		return SFS

	def __iter__(self):
		self.itr = iter(self.sfs)
		return self

	def next(self):
		return self.itr.next()

	def __str__(self):
		data = np.matrix(self.sfs)
		return str(data)
	
	def __len__(self):
		return len(self.sfs)

